# api_server.py
import os
import json
import requests
from typing import List, Dict, Any, Optional, Union
import logging
from pathlib import Path
import tempfile
import shutil
from fastapi import FastAPI, UploadFile, File, HTTPException, Form
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import uvicorn

# 文件处理库
try:
    from docx import Document
except ImportError:
    print("请安装python-docx: pip install python-docx")

try:
    import openpyxl
except ImportError:
    print("请安装openpyxl: pip install openpyxl")

try:
    import pdfplumber
except ImportError:
    print("请安装pdfplumber: pip install pdfplumber")

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 创建FastAPI应用
app = FastAPI(title="AI Image Generator API", description="AI图像生成API服务", version="1.0.0")

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 生产环境中应该限制为具体的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# 请求模型
class OptimizePromptRequest(BaseModel):
    user_requirements: str
    file_content: Optional[str] = None


class GenerateImageRequest(BaseModel):
    prompt: str
    model: str = "gemini-nano-banana"
    size: str = "1024x1024"
    aspect_ratio: str = "1:1"
    quality: str = "medium"
    style_preset: str = "3d-model"
    response_format: str = "url"
    output_format: str = "png"
    n: int = 1
    watermark: bool = False
    negative_prompt: Optional[str] = None
    image: Optional[str] = None
    strength: Optional[float] = None
    mask: Optional[str] = None
    background: Optional[str] = None
    async_mode: bool = False
    retries: int = 0


class WorkflowRequest(BaseModel):
    user_requirements: str
    image_params: Optional[Dict[str, Any]] = {}


# 响应模型
class ApiResponse(BaseModel):
    success: bool
    data: Optional[Dict[str, Any]] = None
    message: str
    error: Optional[str] = None


class AIImageGenerator:
    def __init__(self, api_key: str = None, base_url: str = "https://geekai.co/api/v1"):
        """初始化AI图像生成器"""
        self.api_key = 'sk-z3kAzZ7vYbulURVfWGXBvYFpHfVl9FthWRfUeLaQv9iTNNBr'
        if not self.api_key:
            raise ValueError("请设置GEEKAI_API_KEY环境变量")

        self.base_url = base_url
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

    def read_file_content(self, file_path: str) -> str:
        """读取文件内容"""
        file_path = Path(file_path)
        if not file_path.exists():
            raise FileNotFoundError(f"文件不存在: {file_path}")

        file_extension = file_path.suffix.lower()

        try:
            if file_extension in ['.doc', '.docx']:
                return self._read_word_file(file_path)
            elif file_extension in ['.xls', '.xlsx']:
                return self._read_excel_file(file_path)
            elif file_extension == '.pdf':
                return self._read_pdf_file(file_path)
            else:
                raise ValueError(f"不支持的文件格式: {file_extension}")
        except Exception as e:
            logger.error(f"读取文件失败: {file_path}, 错误: {str(e)}")
            raise

    def _read_word_file(self, file_path: Path) -> str:
        """读取Word文档内容"""
        doc = Document(file_path)
        content = []

        for paragraph in doc.paragraphs:
            if paragraph.text.strip():
                content.append(paragraph.text.strip())

        for table in doc.tables:
            for row in table.rows:
                row_data = [cell.text.strip() for cell in row.cells if cell.text.strip()]
                if row_data:
                    content.append(" | ".join(row_data))

        return "\n".join(content)

    def _read_excel_file(self, file_path: Path) -> str:
        """读取Excel文档内容"""
        workbook = openpyxl.load_workbook(file_path, data_only=True)
        content = []

        for sheet_name in workbook.sheetnames:
            sheet = workbook[sheet_name]
            content.append(f"工作表: {sheet_name}")

            for row in sheet.iter_rows(values_only=True):
                row_data = [str(cell) for cell in row if cell is not None]
                if any(cell.strip() for cell in row_data if isinstance(cell, str)):
                    content.append(" | ".join(row_data))
            content.append("")

        return "\n".join(content)

    def _read_pdf_file(self, file_path: Path) -> str:
        """读取PDF文档内容"""
        content = []

        with pdfplumber.open(file_path) as pdf:
            for page_num, page in enumerate(pdf.pages, 1):
                text = page.extract_text()
                if text:
                    content.append(f"第{page_num}页:")
                    content.append(text)
                    content.append("")

        return "\n".join(content)

    def call_gpt_api(self, messages: List[Dict[str, str]], model: str = "gpt-5-mini") -> str:
        """调用GPT API"""
        url = f"{self.base_url}/chat/completions"
        payload = {
            "model": model,
            "messages": messages,
            "stream": False
        }

        try:
            response = requests.post(url, headers=self.headers, json=payload, timeout=60)
            response.raise_for_status()

            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                return result['choices'][0]['message']['content']
            else:
                raise ValueError("API响应格式错误")
        except Exception as e:
            logger.error(f"GPT API调用失败: {str(e)}")
            raise

    def process_uploaded_files(self, file_paths: List[str]) -> str:
        """处理上传的文件"""
        all_content = []
        for file_path in file_paths:
            try:
                content = self.read_file_content(file_path)
                all_content.append(f"文件: {Path(file_path).name}\n{content}\n")
            except Exception as e:
                logger.warning(f"跳过文件 {file_path}: {str(e)}")
                continue

        if not all_content:
            raise ValueError("没有成功读取任何文件内容")

        combined_content = "\n".join(all_content)

        messages = [
            {
                "role": "system",
                "content": "你是一个专业的内容整理专家。请仔细分析用户上传的文件内容，提取关键信息，整理成结构化的摘要。包括：1. 主要内容概述 2. 关键数据和信息点 3. 重要结论或建议。请保持内容的准确性和完整性。"
            },
            {
                "role": "user",
                "content": f"请帮我整理以下文件内容：\n\n{combined_content}"
            }
        ]

        return self.call_gpt_api(messages)

    def optimize_prompt(self, user_requirements: str, file_content: str = None) -> str:
        """优化提示词"""
        if file_content:
            # 基于文件内容优化
            messages = [
                {
                    "role": "system",
                    "content": "你是一个专业的图像描述和创意文案专家。基于已整理的文档内容和用户需求，生成详细、准确、富有创意的图像生成提示词。提示词应该包括：1. 主要视觉元素 2. 风格和色调 3. 构图和布局 4. 细节描述。确保提示词能够生成高质量、符合需求的图像。"
                },
                {
                    "role": "user",
                    "content": f"基于以下整理的文档内容：\n\n{file_content}\n\n用户需求：{user_requirements}\n\n请生成优化的图像生成提示词。"
                }
            ]
        else:
            # 直接优化用户需求
            messages = [
                {
                    "role": "system",
                    "content": "你是一个专业的图像描述和创意文案专家。基于用户的需求描述，生成详细、准确、富有创意的图像生成提示词。你需要：1. 分析用户需求的核心要素 2. 补充必要的视觉细节描述 3. 优化语言表达使其更适合AI图像生成 4. 确保提示词具体、清晰且富有表现力。提示词应该包括主要视觉元素、风格色调、构图布局、光影效果等细节。"
                },
                {
                    "role": "user",
                    "content": f"用户需求：{user_requirements}\n\n请将这个需求转换并优化为详细的图像生成提示词。"
                }
            ]

        return self.call_gpt_api(messages)

    def generate_image(self, prompt: str, **kwargs) -> Dict[str, Any]:
        """生成图片"""
        url = f"{self.base_url}/images/generations"

        payload = {
            "model": kwargs.get("model", "gemini-nano-banana"),
            "prompt": prompt,
            "size": kwargs.get("size", "1024x1024"),
            "aspect_ratio": kwargs.get("aspect_ratio", "1:1"),
            "n": kwargs.get("n", 1),
            "quality": kwargs.get("quality", "medium"),
            "style_preset": kwargs.get("style_preset", "3d-model"),
            "response_format": kwargs.get("response_format", "url"),
            "output_format": kwargs.get("output_format", "png"),
            "watermark": kwargs.get("watermark", False),
            "async": kwargs.get("async_mode", False),
            "retries": kwargs.get("retries", 0)
        }

        # 添加可选参数
        optional_params = ["negative_prompt", "image", "strength", "mask", "background"]
        for param in optional_params:
            if param in kwargs and kwargs[param] is not None:
                payload[param] = kwargs[param]

        try:
            response = requests.post(url, headers=self.headers, json=payload, timeout=120)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"图片生成API调用失败: {str(e)}")
            raise


# 全局AI生成器实例
ai_generator = AIImageGenerator()


@app.post("/api/upload-files", response_model=ApiResponse)
async def upload_files(files: List[UploadFile] = File(...)):
    """上传并处理文件"""
    try:
        temp_dir = tempfile.mkdtemp()
        file_paths = []

        # 保存上传的文件
        for file in files:
            if file.filename:
                file_path = Path(temp_dir) / file.filename
                with open(file_path, "wb") as buffer:
                    shutil.copyfileobj(file.file, buffer)
                file_paths.append(str(file_path))

        # 处理文件内容
        organized_content = ai_generator.process_uploaded_files(file_paths)

        # 清理临时文件
        shutil.rmtree(temp_dir, ignore_errors=True)

        return ApiResponse(
            success=True,
            data={"organized_content": organized_content},
            message="文件处理成功"
        )

    except Exception as e:
        # 清理临时文件
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)

        logger.error(f"文件处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/optimize-prompt", response_model=ApiResponse)
async def optimize_prompt(request: OptimizePromptRequest):
    """优化提示词"""
    try:
        optimized_prompt = ai_generator.optimize_prompt(
            user_requirements=request.user_requirements,
            file_content=request.file_content
        )

        return ApiResponse(
            success=True,
            data={"optimized_prompt": optimized_prompt},
            message="提示词优化成功"
        )

    except Exception as e:
        logger.error(f"提示词优化失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/generate-image", response_model=ApiResponse)
async def generate_image(request: GenerateImageRequest):
    """生成图片"""
    try:
        image_result = ai_generator.generate_image(
            prompt=request.prompt,
            model=request.model,
            size=request.size,
            aspect_ratio=request.aspect_ratio,
            quality=request.quality,
            style_preset=request.style_preset,
            response_format=request.response_format,
            output_format=request.output_format,
            n=request.n,
            watermark=request.watermark,
            negative_prompt=request.negative_prompt,
            image=request.image,
            strength=request.strength,
            mask=request.mask,
            background=request.background,
            async_mode=request.async_mode,
            retries=request.retries
        )

        return ApiResponse(
            success=True,
            data={"image_result": image_result},
            message="图片生成成功"
        )

    except Exception as e:
        logger.error(f"图片生成失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/api/process-workflow", response_model=ApiResponse)
async def process_workflow(
        user_requirements: str = Form(...),
        files: Optional[List[UploadFile]] = File(None),
        image_params: str = Form("{}")
):
    """完整的工作流程"""
    try:
        # 解析图片参数
        image_params_dict = json.loads(image_params) if image_params else {}

        organized_content = None
        temp_dir = None

        # 如果有文件，先处理文件
        if files and len(files) > 0:
            temp_dir = tempfile.mkdtemp()
            file_paths = []

            for file in files:
                if file.filename:
                    file_path = Path(temp_dir) / file.filename
                    with open(file_path, "wb") as buffer:
                        shutil.copyfileobj(file.file, buffer)
                    file_paths.append(str(file_path))

            # 处理文件内容
            organized_content = ai_generator.process_uploaded_files(file_paths)

        # 优化提示词
        optimized_prompt = ai_generator.optimize_prompt(
            user_requirements=user_requirements,
            file_content=organized_content
        )

        # 生成图片
        image_result = ai_generator.generate_image(
            prompt=optimized_prompt,
            **image_params_dict
        )

        # 清理临时文件
        if temp_dir:
            shutil.rmtree(temp_dir, ignore_errors=True)

        result_data = {
            "workflow_type": "with_files" if organized_content else "direct",
            "organized_content": organized_content,
            "optimized_prompt": optimized_prompt,
            "image_result": image_result
        }

        return ApiResponse(
            success=True,
            data=result_data,
            message="完整工作流程执行成功"
        )

    except Exception as e:
        # 清理临时文件
        if temp_dir:
            shutil.rmtree(temp_dir, ignore_errors=True)

        logger.error(f"工作流程执行失败: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/api/health")
async def health_check():
    """健康检查"""
    return {"status": "healthy", "message": "AI Image Generator API is running"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8807)
